package org.apache.lucene.search;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.util.PriorityQueue;

import java.io.IOException;
import java.text.Collator;
import java.util.Locale;

/**
 * Expert: A hit queue for sorting by hits by terms in more than one field. Uses
 * <code>FieldCache.DEFAULT</code> for maintaining internal term lookup
 * tables.
 * 
 * <p>
 * Created: Dec 8, 2003 12:56:03 PM
 * 
 * @author Tim Jones (Nacimiento Software)
 * @since lucene 1.4
 * @version $Id: FieldSortedHitQueue.java 605225 2007-12-18 15:13:05Z gsingers $
 * @see Searcher#search(Query,Filter,int,Sort)
 * @see FieldCache
 */
public class FieldSortedHitQueue extends PriorityQueue {

	/**
	 * Creates a hit queue sorted by the given list of fields.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fields
	 *            Fieldable names, in priority order (highest priority first).
	 *            Cannot be <code>null</code> or empty.
	 * @param size
	 *            The number of hits to retain. Must be greater than zero.
	 * @throws IOException
	 */
	public FieldSortedHitQueue(IndexReader reader, SortField[] fields, int size)
			throws IOException {
		final int n = fields.length;
		comparators = new ScoreDocComparator[n];
		this.fields = new SortField[n];
		for (int i = 0; i < n; ++i) {
			String fieldname = fields[i].getField();
			comparators[i] = getCachedComparator(reader, fieldname, fields[i]
					.getType(), fields[i].getLocale(), fields[i].getFactory());

			if (comparators[i].sortType() == SortField.STRING) {
				this.fields[i] = new SortField(fieldname,
						fields[i].getLocale(), fields[i].getReverse());
			} else {
				this.fields[i] = new SortField(fieldname, comparators[i]
						.sortType(), fields[i].getReverse());
			}
		}
		initialize(size);
	}

	/** Stores a comparator corresponding to each field being sorted by */
	protected ScoreDocComparator[] comparators;

	/** Stores the sort criteria being used. */
	protected SortField[] fields;

	/** Stores the maximum score value encountered, needed for normalizing. */
	protected float maxscore = Float.NEGATIVE_INFINITY;

	/**
	 * returns the maximum score encountered by elements inserted via insert()
	 */
	public float getMaxScore() {
		return maxscore;
	}

	// Update maxscore.
	private final void updateMaxScore(FieldDoc fdoc) {
		maxscore = Math.max(maxscore, fdoc.score);
	}

	// The signature of this method takes a FieldDoc in order to avoid
	// the unneeded cast to retrieve the score.
	// inherit javadoc
	public boolean insert(FieldDoc fdoc) {
		updateMaxScore(fdoc);
		return super.insert(fdoc);
	}

	// This overrides PriorityQueue.insert() so that insert(FieldDoc) that
	// keeps track of the score isn't accidentally bypassed.
	// inherit javadoc
	public boolean insert(Object fdoc) {
		return insert((FieldDoc) fdoc);
	}

	// This overrides PriorityQueue.insertWithOverflow() so that
	// updateMaxScore(FieldDoc) that keeps track of the score isn't accidentally
	// bypassed.
	public Object insertWithOverflow(Object element) {
		updateMaxScore((FieldDoc) element);
		return super.insertWithOverflow(element);
	}

	/**
	 * Returns whether <code>a</code> is less relevant than <code>b</code>.
	 * 
	 * @param a
	 *            ScoreDoc
	 * @param b
	 *            ScoreDoc
	 * @return <code>true</code> if document <code>a</code> should be sorted
	 *         after document <code>b</code>.
	 */
	protected boolean lessThan(final Object a, final Object b) {
		final ScoreDoc docA = (ScoreDoc) a;
		final ScoreDoc docB = (ScoreDoc) b;

		// run comparators
		final int n = comparators.length;
		int c = 0;
		for (int i = 0; i < n && c == 0; ++i) {
			c = (fields[i].reverse) ? comparators[i].compare(docB, docA)
					: comparators[i].compare(docA, docB);
		}
		// avoid random sort order that could lead to duplicates (bug #31241):
		if (c == 0)
			return docA.doc > docB.doc;
		return c > 0;
	}

	/**
	 * Given a FieldDoc object, stores the values used to sort the given
	 * document. These values are not the raw values out of the index, but the
	 * internal representation of them. This is so the given search hit can be
	 * collated by a MultiSearcher with other search hits.
	 * 
	 * @param doc
	 *            The FieldDoc to store sort values into.
	 * @return The same FieldDoc passed in.
	 * @see Searchable#search(Weight,Filter,int,Sort)
	 */
	FieldDoc fillFields(final FieldDoc doc) {
		final int n = comparators.length;
		final Comparable[] fields = new Comparable[n];
		for (int i = 0; i < n; ++i)
			fields[i] = comparators[i].sortValue(doc);
		doc.fields = fields;
		// if (maxscore > 1.0f) doc.score /= maxscore; // normalize scores
		return doc;
	}

	/** Returns the SortFields being used by this hit queue. */
	SortField[] getFields() {
		return fields;
	}

	static ScoreDocComparator getCachedComparator(IndexReader reader,
			String field, int type, Locale locale, SortComparatorSource factory)
			throws IOException {
		if (type == SortField.DOC)
			return ScoreDocComparator.INDEXORDER;
		if (type == SortField.SCORE)
			return ScoreDocComparator.RELEVANCE;
		FieldCacheImpl.Entry entry = (factory != null) ? new FieldCacheImpl.Entry(
				field, factory)
				: new FieldCacheImpl.Entry(field, type, locale);
		return (ScoreDocComparator) Comparators.get(reader, entry);
	}

	/**
	 * Internal cache of comparators. Similar to FieldCache, only caches
	 * comparators instead of term values.
	 */
	static final FieldCacheImpl.Cache Comparators = new FieldCacheImpl.Cache() {

		protected Object createValue(IndexReader reader, Object entryKey)
				throws IOException {
			FieldCacheImpl.Entry entry = (FieldCacheImpl.Entry) entryKey;
			String fieldname = entry.field;
			int type = entry.type;
			Locale locale = entry.locale;
			SortComparatorSource factory = (SortComparatorSource) entry.custom;
			ScoreDocComparator comparator;
			switch (type) {
			case SortField.AUTO:
				comparator = comparatorAuto(reader, fieldname);
				break;
			case SortField.INT:
				comparator = comparatorInt(reader, fieldname);
				break;
			case SortField.FLOAT:
				comparator = comparatorFloat(reader, fieldname);
				break;
			case SortField.LONG:
				comparator = comparatorLong(reader, fieldname);
				break;
			case SortField.DOUBLE:
				comparator = comparatorDouble(reader, fieldname);
				break;
			case SortField.STRING:
				if (locale != null)
					comparator = comparatorStringLocale(reader, fieldname,
							locale);
				else
					comparator = comparatorString(reader, fieldname);
				break;
			case SortField.CUSTOM:
				comparator = factory.newComparator(reader, fieldname);
				break;
			default:
				throw new RuntimeException("unknown field type: " + type);
			}
			return comparator;
		}
	};

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * integers.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg integer values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorInt(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		final int[] fieldOrder = FieldCache.DEFAULT.getInts(reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				final int fi = fieldOrder[i.doc];
				final int fj = fieldOrder[j.doc];
				if (fi < fj)
					return -1;
				if (fi > fj)
					return 1;
				return 0;
			}

			public Comparable sortValue(final ScoreDoc i) {
				return new Integer(fieldOrder[i.doc]);
			}

			public int sortType() {
				return SortField.INT;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * integers.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg integer values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorLong(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		final long[] fieldOrder = ExtendedFieldCache.EXT_DEFAULT.getLongs(
				reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				final long li = fieldOrder[i.doc];
				final long lj = fieldOrder[j.doc];
				if (li < lj)
					return -1;
				if (li > lj)
					return 1;
				return 0;
			}

			public Comparable sortValue(final ScoreDoc i) {
				return new Long(fieldOrder[i.doc]);
			}

			public int sortType() {
				return SortField.LONG;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * floats.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg float values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorFloat(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		final float[] fieldOrder = FieldCache.DEFAULT.getFloats(reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				final float fi = fieldOrder[i.doc];
				final float fj = fieldOrder[j.doc];
				if (fi < fj)
					return -1;
				if (fi > fj)
					return 1;
				return 0;
			}

			public Comparable sortValue(final ScoreDoc i) {
				return new Float(fieldOrder[i.doc]);
			}

			public int sortType() {
				return SortField.FLOAT;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * doubles.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg float values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorDouble(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		final double[] fieldOrder = ExtendedFieldCache.EXT_DEFAULT.getDoubles(
				reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				final double di = fieldOrder[i.doc];
				final double dj = fieldOrder[j.doc];
				if (di < dj)
					return -1;
				if (di > dj)
					return 1;
				return 0;
			}

			public Comparable sortValue(final ScoreDoc i) {
				return new Double(fieldOrder[i.doc]);
			}

			public int sortType() {
				return SortField.DOUBLE;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * strings.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg string values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorString(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		final FieldCache.StringIndex index = FieldCache.DEFAULT.getStringIndex(
				reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				final int fi = index.order[i.doc];
				final int fj = index.order[j.doc];
				if (fi < fj)
					return -1;
				if (fi > fj)
					return 1;
				return 0;
			}

			public Comparable sortValue(final ScoreDoc i) {
				return index.lookup[index.order[i.doc]];
			}

			public int sortType() {
				return SortField.STRING;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to a field containing
	 * strings.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg string values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorStringLocale(final IndexReader reader,
			final String fieldname, final Locale locale) throws IOException {
		final Collator collator = Collator.getInstance(locale);
		final String field = fieldname.intern();
		final String[] index = FieldCache.DEFAULT.getStrings(reader, field);
		return new ScoreDocComparator() {

			public final int compare(final ScoreDoc i, final ScoreDoc j) {
				String is = index[i.doc];
				String js = index[j.doc];
				if (is == js) {
					return 0;
				} else if (is == null) {
					return -1;
				} else if (js == null) {
					return 1;
				} else {
					return collator.compare(is, js);
				}
			}

			public Comparable sortValue(final ScoreDoc i) {
				return index[i.doc];
			}

			public int sortType() {
				return SortField.STRING;
			}
		};
	}

	/**
	 * Returns a comparator for sorting hits according to values in the given
	 * field. The terms in the field are looked at to determine whether they
	 * contain integers, floats or strings. Once the type is determined, one of
	 * the other static methods in this class is called to get the comparator.
	 * 
	 * @param reader
	 *            Index to use.
	 * @param fieldname
	 *            Fieldable containg values.
	 * @return Comparator for sorting hits.
	 * @throws IOException
	 *             If an error occurs reading the index.
	 */
	static ScoreDocComparator comparatorAuto(final IndexReader reader,
			final String fieldname) throws IOException {
		final String field = fieldname.intern();
		Object lookupArray = ExtendedFieldCache.EXT_DEFAULT.getAuto(reader,
				field);
		if (lookupArray instanceof FieldCache.StringIndex) {
			return comparatorString(reader, field);
		} else if (lookupArray instanceof int[]) {
			return comparatorInt(reader, field);
		} else if (lookupArray instanceof long[]) {
			return comparatorLong(reader, field);
		} else if (lookupArray instanceof float[]) {
			return comparatorFloat(reader, field);
		} else if (lookupArray instanceof String[]) {
			return comparatorString(reader, field);
		} else {
			throw new RuntimeException("unknown data type in field '" + field
					+ "'");
		}
	}
}
